﻿using Encrypt;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Primitives;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Notes_CommonDev.MiddleWare
{
    public class EncryptMiddleware
    {
        private readonly RequestDelegate _next;

        //Your constructor will have the dependencies needed for database access
        public EncryptMiddleware(RequestDelegate next)
        {
            _next = next;
        }

        public async Task Invoke(HttpContext context)
        {
            if (!context.Request.Headers.TryGetValue("Encrypt", out StringValues vs) || vs != "1")
            {
                await _next(context);
            }
            else
            {
                var rawRequestBody = context.Request.Body;
                var rawResponseBody = context.Response.Body;
                using (var newRequest = new MemoryStream())
                {
                    // 替换request流
                    context.Request.Body = newRequest;
                    using (var newResponse = new MemoryStream())
                    {
                        // 替换response流
                        context.Response.Body = newResponse;

                        string reqdeStr = null;
                        using (var reader = new StreamReader(rawRequestBody))
                        {
                            newResponse.Seek(0,SeekOrigin.Begin);
                            reqdeStr = await reader.ReadToEndAsync();
                            if (!string.IsNullOrWhiteSpace(reqdeStr))
                            {
                                reqdeStr = AES.AesDecrypt(reqdeStr, AES.SimpleKey);
                            }
                        }

                        using (var writer = new StreamWriter(newRequest))
                        {
                            await writer.WriteAsync(reqdeStr);
                            await writer.FlushAsync();
                            newRequest.Seek(0, SeekOrigin.Begin);
                            context.Request.Body = newRequest;
                            await _next(context);
                        }

                        string resdeStr = null;
                        using (var reader = new StreamReader(newResponse))
                        {
                            newResponse.Seek(0, SeekOrigin.Begin);
                            resdeStr = await reader.ReadToEndAsync();
                            if (!string.IsNullOrWhiteSpace(resdeStr))
                            {
                                resdeStr = AES.AesEncrypt(resdeStr, AES.SimpleKey);
                            }
                        }

                        using (var writer = new StreamWriter(rawResponseBody))
                        {
                            await writer.WriteAsync(resdeStr);
                            await writer.FlushAsync();
                        }
                    }
                }
            }

        }

    }
}
